#!/usr/bin/python3
# -*- coding: utf-8 -*-

import sys
# sys.path.append('/opt/work/caffe/python')
sys.path.insert(0, '.')

import argparse
import os
import struct
import cv2

video_path = r'/rootfs/media/kasim/DataSet/Record2'
video_list_bbox_file = r'/rootfs/media/kasim/DataSet/Record2/detect/video_bbox_count.txt'
video_list_bbox_path = r'/rootfs/media/kasim/DataSet/Record2/detect'

video_list_file = r'/rootfs/media/kasim/DataSet/Record2/file_list.txt'

DISPLAY_WIDTH = 1280
DISPLAY_HEIGHT = 720
MIN_BBOX_COUNT = 20
DISPLAY_ALL_FRAME = True
DISPLAY_CLASS_LIST = [
    'Person',
    'Cat',
    'Dog',
    'BabyCar',
    'Face',
]

SHOW = True
THRESHOLDS = [
    0.5,  # Person
    0.7,  # Cat
    0.7,  # Dog
    0.7,  # BabyCar
    0.6,  # Face
]

SLEEP = 1

CLASS_NAMES = [
    'Person',
    'Cat',
    'Dog',
    'BabyCar',
    'Face',
]

CLASS_LABELS = {
    'Person': 0,
    'Cat': 1,
    'Dog': 2,
    'BabyCar': 3,
    'Face': 4,
}

COLORS = [
    (255, 0, 0),
    (0, 255, 0),
    (0, 255, 255),
    (255, 0, 0),
    (0, 0, 255),
]

DISPLAY_LABLE_SET = set()
for class_name in DISPLAY_CLASS_LIST:
    DISPLAY_LABLE_SET.add(CLASS_LABELS[class_name])


def display_video(video_path, bbox_info_dict):
    cap = cv2.VideoCapture(video_path)

    is_pause = True
    frame_id = 0
    while True:
        grabbed, image = cap.read()

        if not grabbed:
            break

        has_display_lable = False
        bbox_info_list = bbox_info_dict.get(frame_id, None)
        if bbox_info_list is not None:
            for bbox_info in bbox_info_list:
                lable = bbox_info[0]
                if lable not in DISPLAY_LABLE_SET:
                    continue
                has_display_lable = True
                color = COLORS[lable]
                _class_name = CLASS_NAMES[lable]
                bbox = bbox_info[1:]
                label_text = '{} {:.2f}'.format(_class_name, bbox[4])
                cv2.putText(image, label_text, (bbox[0], bbox[1] + 32), cv2.FONT_HERSHEY_COMPLEX, 1.0, color)
                cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, thickness=2)
        frame_id += 1

        if DISPLAY_ALL_FRAME or has_display_lable:
            if (image.shape[0] != DISPLAY_HEIGHT) or (image.shape[1] != DISPLAY_WIDTH):
                image = cv2.resize(image, (DISPLAY_WIDTH, DISPLAY_HEIGHT))
            cv2.imshow('image', image)
            if is_pause:
                key = cv2.waitKey()
            else:
                key = cv2.waitKey(SLEEP)

            if -1 != key:
                if key == ord('p'):
                    is_pause = not is_pause
                elif key == ord('n'):
                    break
                elif key == ord('q'):
                    return False
                else:
                    while key != ord(' '):
                        key = cv2.waitKey()
                        if key == ord('q'):
                            return False

    return True


def main():
    video_bbox_count_dict = {}
    with open(video_list_bbox_file, 'r') as file:
        for video_info in file.readlines():
            video_info = video_info.split(',')
            video_bbox_count_dict[video_info[0].strip()] = int(video_info[1].strip())

    video_list = []
    with open(video_list_file, 'r') as file:
        for line in file.readlines():
            video_name = line.split()[0].strip()
            video_list.append(video_name)

    file_count = 0
    total_file_count = len(video_list)
    for video_name in video_list:
        file_count += 1
        print('{}/{} {}'.format(file_count, total_file_count, video_name))
        bbox_count = video_bbox_count_dict.get(video_name, 0)
        if bbox_count < MIN_BBOX_COUNT:
            continue

        bbox_info_dict = {}
        with open(os.path.join(video_list_bbox_path, os.path.splitext(video_name)[0]+'.dat'), 'rb') as file:
            for j in range(bbox_count):
                bbox_info = file.read(28)
                if len(bbox_info) < 28:
                    break
                bbox_info = struct.unpack('6i1f', bbox_info)
                frame_id = bbox_info[0]
                label = bbox_info[1]
                score = bbox_info[6]
                if score >= THRESHOLDS[label]:
                    if frame_id not in bbox_info_dict.keys():
                        bbox_info_dict[frame_id] = [bbox_info[1:]]
                    else:
                        bbox_info_dict[frame_id].append(bbox_info[1:])

        if len(bbox_info_dict) < 1:
            continue

        if not display_video(os.path.join(video_path, video_name), bbox_info_dict):
            break

    print('Finish!')


if __name__ == '__main__':
    main()
